"""
@Project : KonwledgeDistilling
@Author  : wxj233
@Time    : 2025/11/6 11:57
@Desc    : 
"""
from torch import nn


class Teacher(nn.Module):
    """
    字体识别模型
    """
    def __init__(self):
        super().__init__()

        self.c1 = nn.Sequential(nn.Conv2d(1, 6, (3, 3)),
                                nn.BatchNorm2d(6),
                                nn.ReLU())  # N*6*26*26

        self.c2 = nn.Sequential(nn.Conv2d(6, 12, (3, 3)),
                                nn.BatchNorm2d(12),
                                nn.ReLU())  # N*12*24*24

        self.pool1 = nn.MaxPool2d(2)  # N*12*12*12

        self.c3 = nn.Sequential(nn.Conv2d(12, 24, (3, 3)),
                                nn.BatchNorm2d(24),
                                nn.ReLU())  # N*24*10*10
        self.c4 = nn.Sequential(nn.Conv2d(24, 48, (3, 3)),
                                nn.BatchNorm2d(48),
                                nn.ReLU())  # N*48*8*8
        self.pool2 = nn.MaxPool2d(2)  # N*48*4*4

        self.c5 = nn.Sequential(nn.Conv2d(48, 24, (1, 1)),
                                nn.BatchNorm2d(24),
                                nn.ReLU())  # N*24*4*4

        self.FFN = nn.Sequential(nn.Linear(24*4*4, 12*4*4),
                                 nn.BatchNorm1d(12*4*4),
                                 nn.ReLU(),
                                 nn.Linear(12*4*4, 12*4),
                                 nn.BatchNorm1d(12*4),
                                 nn.ReLU())  # N*(12*4)

        self.out = nn.Sequential(nn.Linear(12*4, 10))

    def forward(self, X):
        """

        :param X: 图像数据,N*28*28
        :return:
        """
        X = self.c1(X)  # N*6*26*26
        X = self.c2(X)  # N*12*24*24
        X = self.pool1(X)  # N*12*12*12
        X = self.c3(X)  # N*24*10*10
        X = self.c4(X)  # N*48*8*8
        X = self.pool2(X)  # N*48*4*4
        X = self.c5(X)  # N*24*4*4
        X = self.FFN(X.view(-1, 24*4*4))  # N*(12*4)
        X = self.out(X)  # N*10

        return X



class Student(nn.Module):
    """
    字体识别模型
    """
    def __init__(self):
        super().__init__()

        self.c1 = nn.Sequential(nn.Conv2d(1, 6, (5, 5)),
                                nn.BatchNorm2d(6),
                                nn.ReLU())  # N*6*24*24
        self.pool1 = nn.MaxPool2d(2)  # N*6*12*12

        self.c3 = nn.Sequential(nn.Conv2d(6, 12, (3, 3)),
                                nn.BatchNorm2d(12),
                                nn.ReLU())  # N*12*10*10
        self.pool2 = nn.MaxPool2d(2)  # N*12*5*5

        self.FFN = nn.Sequential(nn.Linear(12*5*5, 12),
                                 nn.BatchNorm1d(12),
                                 nn.ReLU())  # N*12

        self.out = nn.Sequential(nn.Linear(12, 10))

    def forward(self, X):
        """

        :param X: 图像数据,N*28*28
        :return:
        """
        X = self.c1(X)  # N*6*24*24
        X = self.pool1(X)  # N*6*12*12
        X = self.c3(X)  # N*12*10*10
        X = self.pool2(X)  # N*12*5*5
        X = self.FFN(X.view(-1, 12*5*5))  # N*12
        X = self.out(X)  # N*10

        return X

